import mysql.connector

table_name = 'BaiduImageRecords'
_create_table_sql_statement = f'''
CREATE TABLE `{table_name}` (
    `id` INT NOT NULL AUTO_INCREMENT,
    `identity` VARCHAR(768) NOT NULL,
    `file_name` TEXT NOT NULL,
    `url` TEXT NOT NULL,
    `source_url` TEXT,
    `source_image_url` TEXT,
    `caption` TEXT,
    `hash` CHAR(64) AS (SHA2(CONCAT(`identity`, `file_name`), 256)) STORED NOT NULL,
    `create_time` TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    `modify_time` TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
    PRIMARY KEY (`id`),
    UNIQUE KEY `unique_hash` (`hash`),
    INDEX `index_identity` (`identity`)
) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci
'''
_drop_table_sql_statement = f'DROP TABLE `{table_name}`'
_exists_file_sql_statement = f'SELECT EXISTS(SELECT * FROM `{table_name}` WHERE `hash` = (SHA2(CONCAT(%s, %s), 256)))'
_new_record_sql_statement = f'''
INSERT INTO `{table_name}` (
    `identity`, `file_name`, `url`, `source_url`, `source_image_url`, `caption`) values (%s, %s, %s, %s, %s, %s)'''
_count_all_sql_statement = f'SELECT COUNT(*) FROM `{table_name}`'
_count_by_identity_sql_statement = f"SELECT COUNT(*) FROM `{table_name}` WHERE `identity` = %s"
_select_id_file_url_sql_statement = f'SELECT `id`, `identity`, `file_name`, `url`, `source_url`, `source_image_url`, `caption` from `{table_name}`'


class BaiduCrawlerDAO:
    def __init__(self, connection_config: dict):
        self.connection_config = connection_config

    def __enter__(self):
        self.ctx = mysql.connector.connect(**self.connection_config)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.ctx.close()

    def get_cursor(self, buffered=False):
        return self.ctx.cursor(buffered=buffered)

    def exists(self, cursor, identity: str, file_name: str):
        cursor.execute(_exists_file_sql_statement, (identity, file_name))
        result = cursor.fetchone()[0]
        return result == 1

    def insert_and_commit(self, cursor, identity: str, file_name: str, url: str, source_url: str, source_image_url: str, caption: str):
        try:
            self.insert(cursor, identity, file_name, url, source_url, source_image_url, caption)
            self.ctx.commit()
            return True, None, None
        except mysql.connector.Error as e:
            self.ctx.rollback()
            return False, e.errno, str(e)

    def insert(self, cursor, identity: str, file_name: str, url: str, source_url: str, source_image_url: str, caption: str):
        cursor.execute(_new_record_sql_statement, (identity, file_name, url, source_url, source_image_url, caption))

    def insert_multiple(self, cursor, identities: list, file_names: list, urls: list, source_urls: list, source_image_urls: list, captions: list):
        assert len(identities) == len(file_names) == len(urls)
        cursor.executemany(_new_record_sql_statement, zip(identities, file_names, urls, source_urls, source_image_urls, captions))

    def commit(self):
        self.ctx.commit()

    def rollback(self):
        self.ctx.rollback()

    def create_table(self, cursor):
        cursor.execute(_create_table_sql_statement)

    def drop_table(self, cursor):
        cursor.execute(_drop_table_sql_statement)

    def count_all(self, cursor):
        cursor.execute(_count_all_sql_statement)
        return cursor.fetchone()[0]

    def count_by_identity(self, cursor, identity):
        cursor.execute(_count_by_identity_sql_statement, (identity,))
        return cursor.fetchone()[0]

    def get_iterator(self, cursor):
        cursor.execute(_select_id_file_url_sql_statement)
        return cursor

    def get_iterator_with_id_limits(self, cursor, id_min: int=None, id_max: int=None):
        if id_min is None and id_max is None:
            cursor.execute(_select_id_file_url_sql_statement)
        elif id_min is not None and id_max is not None:
            cursor.execute(_select_id_file_url_sql_statement + ' WHERE `id` >= %s AND `id` <= %s', (id_min, id_max))
        elif id_min is not None:
            cursor.execute(_select_id_file_url_sql_statement + ' WHERE `id` >= %s', (id_min,))
        elif id_max is not None:
            cursor.execute(_select_id_file_url_sql_statement + ' WHERE `id` <= %s', (id_max,))
        else:
            raise Exception
        return cursor
